module mxu_psb (
    input logic              clk,
    input logic              rst_n,
          psb_rd_req_if.in   mat_rd_req_if,
          psb_rd_req_if.in   aru_rd_req_if,
          psb_wr_req_if.in   psb_wr_req_if,
          mxu_psb_dat_if.out mat_rd_dat_if,
          aru_psb_dat_if.out aru_rd_dat_if

);

    // Parameters
    localparam BANK_WIDTH = `M0 * `N0 * `ACC_Byte;  // Data width per bank in bytes
    localparam BANK_DEPTH = `PSB_SIZE / BANK_WIDTH;  // Number of entries per bank
    localparam ADDR_WIDTH = $clog2(BANK_DEPTH);

    // PSB Storage - organized as M0xN0 fp32 elements per entry
    fp32_t        psb_mem                                                          [BANK_DEPTH-1:0][`M0-1:0][`N0-1:0];

    // ========================================================================
    // Read Arbitration: Prioritize mat_rd_req over aru_rd_req
    // ========================================================================
    logic         arb_grant_mat;  // 1: grant to mat_rd_req, 0: grant to aru_rd_req
    logic         rd_req_vld;
    logic  [23:0] rd_addr;

    // Simple priority arbitration: mat_rd_req has higher priority
    assign arb_grant_mat     = mat_rd_req_if.vld;
    assign rd_req_vld        = mat_rd_req_if.vld || aru_rd_req_if.vld;
    assign rd_addr           = arb_grant_mat ? mat_rd_req_if.psb_addr : aru_rd_req_if.psb_addr;

    // Ready signals for read requesters
    assign mat_rd_req_if.rdy = arb_grant_mat;
    assign aru_rd_req_if.rdy = !mat_rd_req_if.vld && aru_rd_req_if.vld;

    // ========================================================================
    // Read Operation
    // ========================================================================
    logic [ADDR_WIDTH-1:0] rd_idx;
    logic                  rd_req_vld_q;
    logic                  arb_grant_mat_q;

    assign rd_idx = rd_addr[ADDR_WIDTH+$clog2(BANK_WIDTH)-1:$clog2(BANK_WIDTH)];

    // Pipeline stage 1: Register read request
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            rd_req_vld_q    <= 1'b0;
            arb_grant_mat_q <= 1'b0;
        end else begin
            rd_req_vld_q    <= rd_req_vld;
            arb_grant_mat_q <= arb_grant_mat;
        end
    end

    // Read data from memory
    fp32_t rd_data[`M0-1:0][`N0-1:0];
    always_ff @(posedge clk) begin
        if (rd_req_vld) begin
            rd_data <= psb_mem[rd_idx];
        end
    end

    // Output to mat_rd_dat_if (M0xN0)
    assign mat_rd_dat_if.vld = rd_req_vld_q && arb_grant_mat_q;
    assign mat_rd_dat_if.dat = rd_data;
    // rdy 是 input（从外部接收），不能在这里驱动
    // assign mat_rd_dat_if.rdy = 1'b1;  // 删除这行

    // Output to aru_rd_dat_if (P_ARUxN0) - only first P_ARU rows
    assign aru_rd_dat_if.vld = rd_req_vld_q && !arb_grant_mat_q;
    always_comb begin
        for (int i = 0; i < `P_ARU; i++) begin
            aru_rd_dat_if.dat[i] = rd_data[i];
        end
    end
    // rdy 是 input（从外部接收），不能在这里驱动
    // assign aru_rd_dat_if.rdy = 1'b1;  // 删除这行

    // ========================================================================
    // Write Operation
    // ========================================================================
    logic [ADDR_WIDTH-1:0] wr_idx;
    assign wr_idx            = psb_wr_req_if.psb_addr[ADDR_WIDTH+$clog2(BANK_WIDTH)-1:$clog2(BANK_WIDTH)];

    // Write ready - can accept when not conflicting with read
    assign psb_wr_req_if.rdy = 1'b1;

    // Write to memory with msk
    always_ff @(posedge clk) begin
        if (psb_wr_req_if.vld && psb_wr_req_if.rdy) begin
            for (int i = 0; i < `M0; i++) begin
                for (int j = 0; j < `N0; j++) begin
                    if (psb_wr_req_if.msk[i][j]) begin
                        psb_mem[wr_idx][i][j] <= psb_wr_req_if.dat[i][j];
                    end
                end
            end
        end
    end

endmodule

